import requests
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import plotly.graph_objects as go
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error, mean_squared_error, mean_absolute_percentage_error
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense
import tushare as ts
import math
import random
import webbrowser

# 设置随机种子，保证结果的可复现性
random.seed(1)
np.random.seed(1)
tf.random.set_seed(1)

# 设置中文字体为黑体，解决负号显示问题（用于部分和matplotlib相关的操作，比如后续模型训练结果可视化等）
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 获取股票数据的函数，使用Tushare的数据接口（需替换token才能正常获取数据）
def get_stock_data(token, ts_code, start_date, end_date, fields):
    """
    从tushare获取指定股票代码、日期范围以及字段的股票数据

    参数:
    token (str): 在tushare平台注册获取的token
    ts_code (str): 股票代码，如 "01211.HK"
    start_date (str): 开始日期，格式为 "YYYYMMDD"
    end_date (str): 结束日期，格式为 "YYYYMMDD"
    fields (list): 要获取的股票数据字段列表，如 ["ts_code", "trade_date", "open", "high", "low", "close", "pre_close", "change", "pct_chg", "vol", "amount"]

    返回:
    pd.DataFrame: 获取到的股票数据，如果获取失败返回None
    """
    try:
        pro = ts.pro_api(token)  # 请替换为自己的Tushare token
        data = pro.daily(ts_code=ts_code,
                         start_date=start_date,
                         end_date=end_date,
                         fields=fields)
        return data
    except Exception as e:
        print(f"获取股票数据时出现错误: {e}")
        return None


# 数据预处理函数，进行缺失值处理、数据归一化等操作，简化数据处理步骤
def preprocess_data(data):
    df = data.copy()
    # 将交易日期列设置为索引，并转换为日期时间类型
    df.set_index('trade_date', inplace=True)
    df.index = pd.to_datetime(df.index)
    print("设置日期索引后的基本信息：")
    print(df.info())

    # 去除包含缺失值的行
    df.dropna(inplace=True)
    print("去除缺失值后的基本信息：")
    print(df.info())

    # 提取需要的列（开盘价、最高价、最低价、收盘价、成交量）
    columns = ['open', 'high', 'low', 'close', 'vol']
    df = df[columns]
    print("提取关键列后的基本信息：")
    print(df.info())

    # 数据归一化，使用MinMaxScaler将数据归一化到0-1区间
    scaler = MinMaxScaler(feature_range=(0, 1))
    if df.shape[0] == 0:
        print("数据为空，无法进行归一化处理！")
        exit(1)
    df_scaled = scaler.fit_transform(df)
    print("归一化后的数据形状:", df_scaled.shape)
    return df, df_scaled, scaler


# 创建训练数据集和测试数据集的函数，按照时间序列划分数据，避免重复计算相关逻辑优化
def create_dataset(dataset, time_steps=1):
    dataX, dataY = [], []
    for i in range(len(dataset) - time_steps - 1):
        a = dataset[i:(i + time_steps), :]
        dataX.append(a)
        dataY.append(dataset[i + time_steps, -1])
    return np.array(dataX), np.array(dataY)


# 构建并训练多层LSTM模型的函数，简化超参数调优逻辑
def build_and_train_model(X_train, y_train, X_test, y_test):
    units_choices = [32, 50]  # 缩小超参数选择范围
    epochs_choices = [10, 20]
    batch_size_choices = [32, 64]
    best_rmse = float('inf')
    best_model = None
    best_history = None
    for units in units_choices:
        for epochs in epochs_choices:
            for batch_size in batch_size_choices:
                model = Sequential()
                # 多层 LSTM
                model.add(LSTM(units=units, activation='relu', return_sequences=True,
                               input_shape=(X_train.shape[1], X_train.shape[2])))
                model.add(LSTM(units=units, activation='relu'))
                model.add(Dense(units=1))
                model.compile(loss='mean_squared_error', optimizer='adam')
                history = model.fit(X_train, y_train, validation_data=(X_test, y_test), epochs=epochs,
                                    batch_size=batch_size, verbose=1)
                # 在测试集上评估模型，获取RMSE
                y_pred = model.predict(X_test)
                rmse = math.sqrt(mean_squared_error(y_test, y_pred))
                if rmse < best_rmse:
                    best_rmse = rmse
                    best_model = model
                    best_history = history
    return best_model, best_history


# 计算多个评估指标（均方根误差RMSE、平均绝对误差MAE、平均绝对百分比误差MAPE），用于更全面评估模型性能
def calculate_metrics(y_true, y_pred):
    # 确保 y_true 和 y_pred 的长度一致
    min_len = min(len(y_true), len(y_pred))  # 获取较小的长度
    y_true = y_true[:min_len]  # 修剪 y_true
    y_pred = y_pred[:min_len]  # 修剪 y_pred

    rmse = math.sqrt(mean_squared_error(y_true, y_pred))
    mae = mean_absolute_error(y_true, y_pred)
    mape = mean_absolute_percentage_error(y_true, y_pred)
    return rmse, mae, mape


if __name__ == "__main__":
    # Tushare平台获取的token，需替换为你自己的真实token
    token = '572b6125c826e36619a670a9943995412179dabcfb739c5c48913f46'
    ts_code = "601857.SH"  # 股票代码，可替换为其他感兴趣的股票代码
    start_date = "20150101"  # 适当缩小数据获取时间范围，先获取近十年数据进行测试
    end_date = "20241225"
    fields = ["ts_code", "trade_date", "open", "high", "low", "close", "pre_close", "change", "pct_chg", "vol", "amount"]

    # 获取股票数据
    stock_data = get_stock_data(token, ts_code, start_date, end_date, fields)

    # 检查获取的数据是否为空
    if stock_data is None or stock_data.empty:
        print("没有获取到有效的股票数据！")
        exit(1)

    # 数据预处理
    df, df_scaled, scaler = preprocess_data(stock_data)

    # 划分训练集和测试集，使用train_test_split，可指定test_size比例等更灵活
    test_size = 0.2  # 设定测试集占比为20%
    train_data, test_data = train_test_split(df_scaled, test_size=test_size, random_state=1)
    time_steps = 10  # 时间步长
    X_train, y_train = create_dataset(train_data, time_steps)
    X_test, y_test = create_dataset(test_data, time_steps)
    X_train = np.reshape(X_train, (X_train.shape[0], X_train.shape[1], df.shape[1]))
    X_test = np.reshape(X_test, (X_test.shape[0], X_test.shape[1], df.shape[1]))

    model, history = build_and_train_model(X_train, y_train, X_test, y_test)

    # 绘制训练过程中的损失曲线
    plt.plot(history.history['loss'], label='训练损失')
    plt.plot(history.history['val_loss'], label='验证损失')
    plt.title('模型训练损失曲线')
    plt.xlabel('Epoch')
    plt.ylabel('损失值')
    plt.legend()
    plt.show()

    # 使用训练好的模型进行预测
    train_predict = model.predict(X_train)
    train_predict = np.reshape(train_predict, (train_predict.shape[0], 1))

    # 对训练集进行反归一化处理
    train_data_trimmed = train_data[time_steps:len(train_predict) + time_steps, 0:df.shape[1] - 1]
    train_predict = scaler.inverse_transform(np.concatenate((train_data_trimmed, train_predict), axis=1))[:, -1]

    # 使用模型进行测试集预测
    test_predict = model.predict(X_test)
    test_predict = np.reshape(test_predict, (test_predict.shape[0], 1))

    test_data_trimmed = test_data[time_steps:len(test_predict) + time_steps, 0:df.shape[1] - 1]
    test_predict = scaler.inverse_transform(np.concatenate((test_data_trimmed, test_predict), axis=1))[:, -1]

    # 打印测试集的预测均方根误差
    test_rmse = calculate_metrics(df['close'].iloc[len(train_predict) + time_steps:len(train_predict) + len(test_predict) + time_steps - 1].values, test_predict)[0]
    print("测试集均方根误差（RMSE）:", test_rmse)

    # 绘制股票价格走势与预测结果对比图
    fig = go.Figure()

    # 添加真实股价走势（训练集）
    fig.add_trace(go.Scatter(x=df.index[time_steps:len(train_predict) + time_steps - 1],
                             y=df['close'].iloc[time_steps:len(train_predict) + time_steps - 1],
                             mode='lines',
                             name='真实股价（训练集）'))

    # 添加真实股价走势（测试集）
    fig.add_trace(go.Scatter(x=df.index[len(train_predict) + time_steps:len(train_predict) + len(test_predict) + time_steps - 1],
                             y=df['close'].iloc[len(train_predict) + time_steps:len(train_predict) + len(test_predict) + time_steps - 1],
                             mode='lines',
                             name='真实股价（测试集）'))

    # 添加预测股价走势（训练集）
    fig.add_trace(go.Scatter(x=df.index[time_steps:len(train_predict) + time_steps - 1],
                             y=train_predict,
                             mode='lines',
                             name='预测股价（训练集）',
                             line=dict(dash='dash')))

    # 添加预测股价走势（测试集）
    fig.add_trace(go.Scatter(x=df.index[len(train_predict) + time_steps:len(train_predict) + len(test_predict) + time_steps - 1],
                             y=test_predict,
                             mode='lines',
                             name='预测股价（测试集）',
                             line=dict(dash='dash')))

    # 设置图表标题和布局
    fig.update_layout(title='股票真实价格与预测价格对比',
                      xaxis_title='日期',
                      yaxis_title='股价',
                      legend=dict(x=0.8, y=0.9),  # 调整图例位置
                      hovermode='x unified')  # 悬停交互模式

    # 显示图表
    fig.show()